// Copyright Leong Zhin Sheng

package cluster.networks.ib;

import src.*;
import cluster.network_devices.*;
import cluster.task_scheduler.task_generator.*;
import cluster.network_controller.*;
import cluster.dispatcher.*;
import cluster.resource_scheduler.*;
import cluster.task_scheduler.*;
import cluster.network_controller.topo_manager.*;

network tree
{
    parameters:
        // it is suggested that the numMachine should be a multiple of the number of lowest level switch
        int numMachine = default(4);
        int numL1switch = default(2);
        int numGpusperMachine = default(8);
        int numNvlink_port = default(6);
        int switch_delay = default(0);
        int clustertype = default(2);

    submodules:
        machine[numMachine]: Machine 
        {
            parameters:
                numGpus = numGpusperMachine;
                my_index = index;
                @display("i=device/server2;p=80,109");
        }

        topo_manager : TopoManager {}

        dispatcher : Dispatcher {}

        resource_scheduler : ResourceScheduler {}

        task_scheduler : TaskScheduler {
            parameters:
                maxnumGpus = numGpusperMachine * numMachine;
                modelsize = 10000000;
        }

        network_controller : NetworkController {
            parameters:
                clustertype = 2;
        }

        sw[numL1switch]: IBSwitch 
        {
            parameters:
                numSwitchPorts = int(numMachine/numL1switch)*numGpusperMachine + 1;
                @display("p=182,58");
            gates:
                port[int(numMachine/numL1switch)*numGpusperMachine + 1];
        }

        core: IBSwitch 
        {
            parameters:
                numSwitchPorts = numL1switch;
                @display("p=182,58");
            gates:
                port[numL1switch];
        }

        virtual_switch: IB_Virtual_Data_Switch 
        {
            parameters:
                numSwitchPorts = numMachine*numGpusperMachine;
            gates:
                port[numMachine*numGpusperMachine];
        }

    connections allowunconnected:
	    for i=0..numMachine-1, for j=0..numGpusperMachine-1 
	    {
	    virtual_switch.port[i*numGpusperMachine+j] <--> ControlPlaneLink <--> machine[i].virtual_data_switch_port[j];
        machine[i].out_port[j] <--> IB4XQDRWire <--> sw[int(i/int(numMachine/numL1switch))].port[(j+i*numGpusperMachine) % (int(numMachine/numL1switch)*numGpusperMachine)];
        }

        // Connect the L1 switches to core switch as well as central controller
        for i = 0..numL1switch-1
        {
            sw[i].port[int(numMachine/numL1switch)*numGpusperMachine] <--> IB4XQDRWire <--> core.port[i];
            sw[i].central_controller_port <--> ControlPlaneLink <--> network_controller.switch_port++;
        }

        // Connect the core switch to central controller
        network_controller.switch_port++ <--> ControlPlaneLink <--> core.central_controller_port;

        // Connect every machines to its respective L1 switch, central controller and virtual data switch
        for i=0..numMachine-1 {
            dispatcher.machine_port++<-->ControlPlaneLink<-->machine[i].controller_port;
        }
        // Connect the task generator and controller
        task_scheduler.port --> dispatcher.task_generator_port;
}